"""
Core registry system for DIL_Frame components.

This module provides a unified registration framework for all components,
including models, learners, and their subcomponents.
"""

from typing import Dict, Callable, Optional, Set, TypeVar, Generic, Type, cast
import logging
import inspect

# Type variables for generics
I = TypeVar("I")  # Interface type
T = TypeVar("T")  # Concrete type parameter


class Registry(Generic[I]):
    """Central registry for DIL_Frame components.

    A unified registry that can manage any component type that implements
    the specified interface.
    """

    def __init__(self, name: str, interface_type: Type[I]):
        """Initialize a new registry.

        Args:
            name: Name of this registry (e.g. "model", "learner")
            interface_type: Interface type that components must implement
        """
        self.name = name
        self._interface_type = interface_type
        self._registry: Dict[str, Type[I]] = {}
        self._registered_names: Set[str] = set()

    def register(self, name: Optional[str] = None) -> Callable[[Type[I]], Type[I]]:
        """Register a component with the registry.

        Args:
            name: Optional name for registration. If None, uses class name.
                 Names are stored lowercase for case-insensitive lookups.

        Returns:
            Decorator function that registers the decorated class

        Example:
            @MODEL_REGISTRY.register("resnet18")
            class ResNet18(ModelInterface):
                ...
        """

        def _register(obj: Type[I]) -> Type[I]:
            # Validate interface implementation
            if not (inspect.isclass(obj) and issubclass(obj, self._interface_type)):
                raise TypeError(
                    f"Object {obj.__name__} must be a subclass of {self._interface_type.__name__}"
                )

            # Use provided name or class name
            reg_name = name or obj.__name__
            if reg_name in self._registry:
                logging.warning(
                    f"Overriding existing registration for '{reg_name}' in {self.name} registry"
                )

            # Register the component
            self._registry[reg_name.lower()] = obj
            self._registered_names.add(reg_name.lower())

            # Store metadata on the class
            if hasattr(obj, "__registry_info__"):
                registry_info = getattr(obj, "__registry_info__")
                registry_info["registries"].append(self.name)
            else:
                setattr(
                    obj,
                    "__registry_info__",
                    {
                        "name": reg_name.lower(),
                        "registries": [self.name],
                    },
                )

            return obj

        return _register

    def get(self, name: str) -> Type[I]:
        """Get a component class from the registry by name.

        Args:
            name: Name of the component to retrieve

        Returns:
            The registered component class

        Raises:
            KeyError: If component isn't registered
        """
        name = name.lower()
        if name not in self._registry:
            available = sorted(self._registered_names)
            raise KeyError(
                f"Component '{name}' not found in {self.name} registry. "
                f"Available components: {available}"
            )

        return self._registry[name]

    def create(self, name: str, **kwargs) -> I:
        """Create an instance of a registered component.

        Args:
            name: Name of the component to create
            **kwargs: Arguments to pass to the constructor

        Returns:
            An instance of the specified component

        Raises:
            ValueError: If component isn't registered
            Exception: If component creation fails
        """
        try:
            component_class = self.get(name.lower())
            return component_class(**kwargs)
        except KeyError as e:
            available = sorted(self.list_registered())
            raise ValueError(
                f"Unknown {self.name} '{name}'. Available options: {available}"
            ) from e
        except Exception as e:
            logging.error(f"Error creating {self.name} '{name}': {str(e)}")
            raise

    def create_typed(self, name: str, expected_type: Type[T], **kwargs) -> T:
        """Create an instance with preserved concrete type.

        Args:
            name: Name of the component to create
            expected_type: The expected concrete type
            **kwargs: Arguments for the constructor

        Returns:
            A typed instance of the component

        Raises:
            TypeError: If created instance isn't of expected_type
        """
        instance = self.create(name, **kwargs)
        if not isinstance(instance, expected_type):
            raise TypeError(
                f"Created component '{name}' is type {type(instance).__name__}, "
                f"not expected type {expected_type.__name__}"
            )
        return cast(T, instance)

    def list_registered(self) -> Set[str]:
        """List all registered component names."""
        return set(self._registered_names)

    def __contains__(self, name: str) -> bool:
        """Check if a component is registered."""
        return name.lower() in self._registry
